#!/usr/bin/env python3
# -*- coding:utf-8 -*-

"""
【节点名称】：
    RTSPStreamerNode
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
【发布类型】：
    无
"""


import threading
import time
import cv2
import os
import json
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg, BaseNode, get_extra_args
import argparse
from typing import Union
import subprocess


class RTSPStreamerNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )

        self.image_topic = self.get_param("image_topic", "/live/sensor/image_raw")
        self.resize = self.get_param("resize", [640, 480])
        self.fps = self.get_param("fps", 30)
        self.rtsp_url = self.get_param("rtsp_url", "rtsp://127.0.0.1:8554/live")
        self.params_help()

        input_topic = self.image_topic if len(self.image_topic) > 0 else '/' + job_name + '/sensor/image_raw'

        self._image_reader = Subscriber(
            input_topic, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )
        self.image_queue = Queue()
        self.queue_pool.append(self.image_queue)

        width = self.resize[0]
        height = self.resize[1]

        # FFmpeg命令（通过管道输入）
        self.command = [
            'ffmpeg',
            '-y',  # 覆盖输出文件
            '-f', 'rawvideo',
            '-vcodec', 'rawvideo',
            '-pix_fmt', 'bgr24',  # OpenCV默认格式
            '-s', f'{width}x{height}',
            '-r', str(self.fps),
            '-i', '-',  # 从标准输入读取
            '-c:v', 'mpeg4',
            '-pix_fmt', 'yuv420p',
            '-f', 'rtsp',
            self.rtsp_url
        ]
        print("Streaming URL: {}".format(self.rtsp_url))

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()

    def image_callback(self, msg):
        img = sms2cvimg(msg)
        self.image_queue.put({'msg': msg, 'img': img})

    def run(self):
        process = subprocess.Popen(self.command, stdin=subprocess.PIPE)
        while self.is_running():
            img_msg = self.image_queue.get(block=True)
            if img_msg is None:
                break

            img, msg = img_msg['img'], img_msg['msg']
            img = cv2.resize(img, self.resize)
            process.stdin.write(img.tobytes())


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = RTSPStreamerNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()
